import json
import os
import django

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "water_applets.settings")
django.setup()
from models.models import User, CustomerServiceRecords, CustomerServiceFeedbackRecord
from utils.snowflake import snowflake
from utils.stream_std_out_callback_handler import QianMaiStreamingStdOutCallbackHandler

from langchain_openai import ChatOpenAI
from channels.generic.websocket import WebsocketConsumer
from django.core.cache import cache
from django.utils import timezone

from water_applets import settings

os.environ["OPENAI_API_BASE"] = settings.OPENAI_API_BASE
os.environ["OPENAI_API_KEY"] = settings.OPENAI_API_KEY


class ChatConsumer(WebsocketConsumer):
    user = None
    number = ""
    llm = None
    callback = None

    def connect(self):
        token = ""
        headers = self.scope["headers"]
        for key, value in headers:
            if key == b"token":
                token = value.decode("utf-8")
                break
        user_id, token = token.split("::")
        self.user = User.objects.get(pk=user_id)
        key = settings.REDIS_APPLETS_LOGIN_TOKEN_KEY.format(user_id=self.user.id)
        if cache.get(key) == token:
            self.number = str(snowflake())[-6:]
            self.callback = QianMaiStreamingStdOutCallbackHandler()
            self.llm = ChatOpenAI(model='gpt-3.5-turbo',
                                  streaming=True,
                                  callbacks=[self.callback],
                                  max_tokens=1024)
            self.accept()
        else:
            self.close(code=200, reason="Invalid token")

    def disconnect(self, close_code):
        print('@@'*10, close_code)

    def receive(self, text_data):
        question = json.loads(text_data).get("question", "")
        record = CustomerServiceRecords.objects.create(customer_number=self.number, question=question,
                                                       question_user=self.user)
        # 对接AI客服获取答案 https://blog.csdn.net/q506610466/article/details/132790633
        # todo 添加客服反馈记录
        res = self.llm.invoke(question)
        # CustomerServiceFeedbackRecord.objects.create(customer_number=self.number, feedback_user=self.user,
        #                                              feedback_content=res.content)
        record.answer = res.content
        record.answer_time = timezone.now()
        record.save()
        self.send(self.number)
        while self.callback.tokens:  # or self.tokens:
            self.send(self.callback.tokens.pop(0))